import com.baomidou.mybatisplus.annotation.DbType;
import com.baomidou.mybatisplus.annotation.FieldFill;
import com.baomidou.mybatisplus.annotation.IdType;
import com.baomidou.mybatisplus.core.toolkit.StringUtils;
import com.baomidou.mybatisplus.generator.AutoGenerator;
import com.baomidou.mybatisplus.generator.InjectionConfig;
import com.baomidou.mybatisplus.generator.config.*;
import com.baomidou.mybatisplus.generator.config.po.TableField;
import com.baomidou.mybatisplus.generator.config.po.TableFill;
import com.baomidou.mybatisplus.generator.config.po.TableInfo;
import com.baomidou.mybatisplus.generator.config.querys.PostgreSqlQuery;
import com.baomidou.mybatisplus.generator.config.rules.DateType;
import com.baomidou.mybatisplus.generator.config.rules.NamingStrategy;
import com.baomidou.mybatisplus.generator.engine.FreemarkerTemplateEngine;

import java.util.*;

/**
 * @author HongCheng
 * <p>创建于 2021/4/13 15:32</p>
 */
public class Generator {
    // TODO 数据库连接地址
//    static String jdbcUrl = "jdbc:oracle:thin:@localhost:1521:orcl";
    static String jdbcUrl = "jdbc:dm://10.3.2.11:5236?schema=GDONEMAP&ignoreCase=true&columnNameUpperCase=false";
    // TODO 数据库连接用户名
    static String username = "GDONEMAP";
    // TODO 数据库连接密码
    static String password = "GDONEMAP";
    // TODO 要生成的表名
    static String tableName = "DIAGNOSIS_PROBLEM_SCENE";
    // TODO 作者
    static String author = "lzy";
    // TODO 输出文件目录
    static String outDir = "F:\\tem\\";

    // TODO 是否使用标准doc注解
    static boolean isStandardDoc = true;
    static String customAnnotationAuthor = "@作者：";
    static String customAnnotationDate = "@日期：";
    static String customAnnotationParam = "@入参：";
    static String customAnnotationReturn = "@出参：";
    // TODO entity 包名
    static String entityPackageName = "com.guodi.model.entity";
    // TODO mapper 包名
    static String mapperPackageName = "com.guodi.model.mapper";
    // TODO xml 包名
    static String xmlPackageName = "com.guodi.model.mapper";
    // TODO service 包名
    static String servicePackageName = "com.guodi.model.service";
    // TODO controller 包名
    static String controllerPackageName = "com.guodi.model.controller";

    // serviceImpl 包名
    static String serviceImplPackageName = servicePackageName + ".impl";
    // 模板主目录
    static String templatePath = "/template";
    // controller 模板地址
    static String controllerTemplate = templatePath + "/" + "controller.java";
    // service 模板地址
    static String serviceTemplate = templatePath + "/" + "service.java";
    // serviceImpl 模板地址
    static String serviceImplTemplate = templatePath + "/" + "serviceImpl.java";
    // entity 模板地址
    static String entityTemplate = templatePath + "/" + "entity.java";
    // mapper 模板地址
    static String mapperTemplate = templatePath + "/" + "mapper.java";
    // mapper xml 模板地址
    static String xmlTemplate = templatePath + "/" + "mapper.xml";

    // oracle驱动
    static String ORACLE_DRIVER = "oracle.jdbc.OracleDriver";
    // postgresql驱动
    static String POSTGRESQL_DRIVER = "org.postgresql.Driver";
    // sqlserver驱动
    static String SQL_SERVER_DRIVER = "com.microsoft.sqlserver.jdbc.SQLServerDriver";
    // mysql驱动
    static String MYSQL_DRIVER = "com.mysql.jdbc.Driver";
    // DM
    static String DM_DRIVER = "dm.jdbc.driver.DmDriver";

    public static void main(String[] args) {

        // 代码生成器
        AutoGenerator mpg = new AutoGenerator();

        // 全局配置
        GlobalConfig gc = new GlobalConfig();
        gc.setOutputDir( outDir + tableName);  // 输出路径
        gc.setFileOverride(false);  // 不覆盖已有文件
        gc.setAuthor(author);  // 作者
        gc.setOpen(false);  // 打开文件夹
        gc.setSwagger2(true); // 实体属性 Swagger2 注解
        gc.setBaseResultMap(true);  // 生成BaseResultMap
        gc.setBaseColumnList(true);  // 生成BaseColumnList
        gc.setDateType(DateType.ONLY_DATE);  // 只使用 java.util.date 代替
        gc.setActiveRecord(false);  // 不开启 ActiveRecord 模式
        gc.setEntityName(null);  //  实体命名方式不进行处理，如果需要处理可写成：%sEntity 生成 UserEntity
        gc.setMapperName("%sMapper");     // mapper 命名方式   例如：%sMapper 生成 UserMapper.java
        gc.setXmlName("%sMapper");    // Mapper xml 命名方式    例如：%sMapper 生成 UserMapper.xml
        gc.setServiceName("I%sService");   // service 命名方式   例如：I%sService 生成 IUserService.java
        gc.setServiceImplName("%sServiceImpl"); // service impl 命名方式   例如：%sServiceImpl 生成 UserServiceImpl
        gc.setControllerName("%sController");  // controller 命名方式   例如：%sController 生成 UserController
        gc.setIdType(IdType.ASSIGN_UUID);  // 主键id使用UUID
        mpg.setGlobalConfig(gc);

        // 数据源配置
        DataSourceConfig dsc = new DataSourceConfig();
        dsc.setUrl(jdbcUrl);
        if (jdbcUrl.toLowerCase().contains("postgresql")){
            dsc.setDriverName(POSTGRESQL_DRIVER);
            // 添加自定义字段
            dsc.setDbQuery(new CustomPostgreSqlQuery());
        }else if (jdbcUrl.toLowerCase().contains("sqlserver")){
            dsc.setDriverName(SQL_SERVER_DRIVER);
        }else if (jdbcUrl.toLowerCase().contains("oracle")){
            dsc.setDriverName(ORACLE_DRIVER);
        }else if (jdbcUrl.toLowerCase().contains("mysql")){
            dsc.setDriverName(MYSQL_DRIVER);
        }else if (jdbcUrl.toLowerCase().contains("dm")){
            dsc.setDriverName(DM_DRIVER);
        }else{
            throw new RuntimeException("数据库地址错误");
        }
        dsc.setUsername(username);
        dsc.setPassword(password);
        mpg.setDataSource(dsc);

        // 数据库表配置
        StrategyConfig strategy = new StrategyConfig();
        strategy.setSkipView(true);        // 是否跳过视图
        strategy.setNaming(NamingStrategy.underline_to_camel);      // 表命名下划线转驼峰
        strategy.setColumnNaming(NamingStrategy.underline_to_camel);  // 列命名下划线转驼峰
        strategy.setTablePrefix(null);  // 表前缀
        strategy.setFieldPrefix(null);  // 字段前缀
        strategy.setSuperControllerClass("com.guodi.core.boot.ctrl.BaseController"); // 控制器父类
        strategy.setChainModel(true);   // 【实体】使用链式模型
        strategy.setEntityLombokModel(true);
        strategy.setEntityBooleanColumnRemoveIsPrefix(false);  // Boolean类型字段不移除is前缀
        strategy.setRestControllerStyle(true);   // 生成 @RestController 控制器
        strategy.setEntityTableFieldAnnotationEnable(true);   // 是否生成实体时，生成字段注解
        List<TableFill> tableFills = new LinkedList<>();
        new TableFill("modify_time", FieldFill.INSERT);
        new TableFill("modifier", FieldFill.INSERT);
        new TableFill("create_time", FieldFill.UPDATE);
        new TableFill("creator", FieldFill.UPDATE);
        strategy.setTableFillList(tableFills);      // 字段填充
        strategy.setControllerMappingHyphenStyle(true);  // @RequestMapping("/managerUserActionHistory") -> @RequestMapping("/manager-user-action-history")
        strategy.setInclude(tableName);         // 要生成的表名
        mpg.setStrategy(strategy);

        // 包配置
        PackageConfig pc = new PackageConfig();
        pc.setModuleName(null);  //父包模块名
        pc.setParent(null);  // 父包名。如果为空，将下面子包名必须写全部， 否则就只需写子包名
        pc.setEntity(entityPackageName); // Entity包名
        pc.setMapper(mapperPackageName); // Mapper包名
        pc.setService(servicePackageName); // Service包名
        pc.setServiceImpl(serviceImplPackageName); // Service Impl包名
        pc.setController(controllerPackageName);  // Controller包名
        pc.setXml(xmlPackageName); // Mapper XML包名
        mpg.setPackageInfo(pc);

        // 模板配置
        TemplateConfig templateConfig = new TemplateConfig();
        templateConfig.setController(controllerTemplate);
        templateConfig.setServiceImpl(serviceImplTemplate);
        templateConfig.setService(serviceTemplate);
        templateConfig.setEntity(entityTemplate);
        templateConfig.setMapper(mapperTemplate);
        templateConfig.setXml(xmlTemplate);
        mpg.setTemplate(templateConfig);


        // 引擎设置
        mpg.setTemplateEngine(new FreemarkerTemplateEngine());

        // 自定义map
        CustomInjectionConfigBuilder customInjectionConfigBuilder = new CustomInjectionConfigBuilder();
        customInjectionConfigBuilder.setDataSourceConfig(dsc);
        if (isStandardDoc == false){
            customInjectionConfigBuilder.setAnnotationAuthor(customAnnotationAuthor);
            customInjectionConfigBuilder.setAnnotationDate(customAnnotationDate);
            customInjectionConfigBuilder.setAnnotationParam(customAnnotationParam);
            customInjectionConfigBuilder.setAnnotationReturn(customAnnotationReturn);
        }
        mpg.setCfg(customInjectionConfigBuilder.build());

        // 执行
        mpg.execute();
    }


    /**
     * 自定义数据库字段查询，主要是增加字段非空判断
     * */
    static class CustomPostgreSqlQuery extends PostgreSqlQuery{
        @Override
        public String tableFieldsSql() {
            return "SELECT A.attnotnull, A.attname AS name,format_type (A.atttypid,A.atttypmod) AS type,col_description (A.attrelid,A.attnum) AS comment, (CASE WHEN (SELECT COUNT (*) FROM pg_constraint AS PC WHERE A.attnum = PC.conkey[1] AND PC.contype = 'p') > 0 THEN 'PRI' ELSE '' END) AS key  FROM pg_class AS C,pg_attribute AS A WHERE A.attrelid='%s.%s'::regclass AND A.attrelid= C.oid AND A.attnum> 0 AND NOT A.attisdropped ORDER  BY A.attnum";
        }
        @Override
        public String[] fieldCustom() {
            return new String[]{"attnotnull"};
        }
    }

    static class CustomInjectionConfigBuilder{
        private CustomInjectionConfig customInjectionConfig;

        public void setDataSourceConfig(DataSourceConfig dsc) {
            init();
            this.customInjectionConfig.dsc = dsc;
        }
        public void setAnnotationAuthor(String annotationAuthor) {
            init();
            this.customInjectionConfig.annotationAuthor = annotationAuthor;
        }
        public void setAnnotationDate(String annotationDate) {
            init();
            this.customInjectionConfig.annotationDate = annotationDate;
        }
        public void setAnnotationParam(String annotationParam) {
            init();
            this.customInjectionConfig.annotationParam = annotationParam;
        }
        public void setAnnotationReturn(String annotationReturn) {
            init();
            this.customInjectionConfig.annotationReturn = annotationReturn;
        }

        public CustomInjectionConfig build(){
            return this.customInjectionConfig;
        }

        private void init(){
            if (Objects.isNull(this.customInjectionConfig)){
                this.customInjectionConfig = new CustomInjectionConfig();
            }
        }
    }

    /**
     * 自定义数据map构建
     * */
    static class CustomInjectionConfig extends InjectionConfig{
        // 数据源配置
        private DataSourceConfig dsc;
        // doc 注释
        private String annotationAuthor = "@author";
        private String annotationDate = "@date";
        private String annotationParam = "@param";
        private String annotationReturn = "@return";

        public CustomInjectionConfig() {}

        public CustomInjectionConfig(DataSourceConfig dsc){
            this.dsc = dsc;
        }

        @Override
        public void initMap() {
            List<TableInfo> tableInfoList = this.getConfig().getTableInfoList();
            Map<String, Object> tableFieldCustomMap = new HashMap<>();
            // doc注释
            tableFieldCustomMap.put("annotationAuthor", annotationAuthor);
            tableFieldCustomMap.put("annotationDate", annotationDate);
            tableFieldCustomMap.put("annotationParam", annotationParam);
            tableFieldCustomMap.put("annotationReturn", annotationReturn);
            // 处理数据库字段
            for (TableInfo tableInfo : tableInfoList) {
                List<TableField> fields = tableInfo.getFields();
                Map<String, Object> fieldCustomMap = new HashMap<>();

                // 对象的驼峰写法
                fieldCustomMap.put("firstLowerServiceName" ,StringUtils.firstToLowerCase(tableInfo.getServiceName()));
                fieldCustomMap.put("firstLowerControllerName" ,StringUtils.firstToLowerCase(tableInfo.getControllerName()));
                fieldCustomMap.put("firstLowerEntityName" ,StringUtils.firstToLowerCase(tableInfo.getEntityName()));
                fieldCustomMap.put("firstLowerMapperName" ,StringUtils.firstToLowerCase(tableInfo.getMapperName()));

                // 处理字段的长度校验和非空校验
                for (TableField field : fields) {
                    if (dsc.getDbType().equals(DbType.POSTGRE_SQL)){
                        // 字符类型的长度
                        if (Objects.nonNull(field.getColumnType().getType()) && Objects.equals(field.getColumnType().getType().toLowerCase(), "string")){
                            String type = field.getType();
                            if (type.contains("(") && type.contains(")")){
                                String length = type.substring(type.indexOf("(") + 1, type.indexOf(")"));
                                fieldCustomMap.put(field.getPropertyName() + "-Length", length);
                            }
                        }
                        if (field.getCustomMap() != null){
                            Object attnotnull = field.getCustomMap().get("attnotnull");
                            if (attnotnull != null && (attnotnull instanceof Boolean && (Boolean)attnotnull
                                    || Objects.equals(attnotnull.toString().toLowerCase(), "t") || Objects.equals(attnotnull.toString().toLowerCase(), "true"))){
                                if (Objects.nonNull(field.getColumnType().getType()) && Objects.equals(field.getColumnType().getType().toLowerCase(), "string")){
                                    fieldCustomMap.put(field.getPropertyName() + "-NotBlank", true);
                                }else{
                                    fieldCustomMap.put(field.getPropertyName() + "-NotNull", true);
                                }
                            }
                        }
                    }
                }
                tableFieldCustomMap.put(tableInfo.getName(), fieldCustomMap);
            }

            this.setMap(tableFieldCustomMap);
        }
    }
}
